import streamlit as st
from pathlib import Path
import logging
from typing import List, Dict

from core.api_client import OllamaClient
from core.state_manager import StateManager
from core.server_manager import ServerManager
from utils.error_handler import ErrorHandler
from ui.chat import ChatInterface

def setup_logging():
    """配置日志系统"""
    log_dir = Path("logs")
    log_dir.mkdir(exist_ok=True)
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler(log_dir / 'app.log')
        ]
    )
    return logging.getLogger(__name__)

def show_connection_page(server_manager: ServerManager):
    """显示服务器连接页面"""
    st.title("连接到Ollama服务器")
    
    # 在session state中保存选中的服务器信息
    if 'selected_server' not in st.session_state:
        st.session_state.selected_server = {
            'host': 'localhost',
            'port': 11434
        }
    
    # 创建两列布局
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("手动连接")
        # 使用session state中的值作为默认值
        host = st.text_input(
            "服务器地址", 
            value=st.session_state.selected_server['host']
        )
        port = st.number_input(
            "端口", 
            value=st.session_state.selected_server['port'],
            min_value=1,
            max_value=65535
        )
        
        if st.button("连接", key="manual_connect"):
            with st.spinner("正在连接..."):
                try:
                    if server_manager.connect(host, port):
                        st.success("连接成功！")
                        st.rerun()
                    else:
                        st.error("连接失败，请检查服务器地址和端口")
                except Exception as e:
                    st.error(f"连接错误: {str(e)}")
    
    with col2:
        st.subheader("扫描网络")
        
        # 添加自定义网段
        custom_network = st.text_input(
            "添加自定义网段 (例如: 192.168.1.0/24)",
            help="输入要扫描的网段，格式为 IP/掩码"
        )
        
        if custom_network:
            if st.button("添加网段"):
                if server_manager.add_network(custom_network):
                    st.success(f"已添加网段: {custom_network}")
                else:
                    st.error("网段格式错误")
        
        # 显示当前网段列表
        networks = server_manager.get_networks()
        if networks:
            st.caption("当前扫描网段:")
            for network in networks:
                cols = st.columns([3, 1])
                cols[0].text(network)
                if cols[1].button("删除", key=f"del_{network}"):
                    server_manager.remove_network(network)
                    st.rerun()
        
        # 扫描按钮
        scan_col1, scan_col2 = st.columns([3, 1])
        with scan_col1:
            scan_button = st.button("开始扫描", key="start_scan")
        with scan_col2:
            clear_button = st.button("清除结果", key="clear_results")
        
        # 保存扫描结果到session state
        if 'scan_results' not in st.session_state:
            st.session_state.scan_results = []
        
        if scan_button:
            progress_text = st.empty()
            progress_bar = st.progress(0)
            
            def update_progress(msg: str):
                progress_text.text(msg)
            
            def update_bar(value: float):
                progress_bar.progress(value)
            
            # 执行扫描
            with st.spinner("正在扫描网络..."):
                servers = server_manager.scan_network(update_progress, update_bar)
                st.session_state.scan_results = servers
        
        if clear_button:
            st.session_state.scan_results = []
            st.rerun()
        
        # 显示扫描结果
        if st.session_state.scan_results:
            st.success(f"发现 {len(st.session_state.scan_results)} 个服务器")
            
            # 使用radio选择服务器
            server_options = [
                f"服务器 {idx + 1}: {server['host']}:{server['port']} (版本: {server.get('version', '未知')})"
                for idx, server in enumerate(st.session_state.scan_results)
            ]
            
            selected_idx = st.radio(
                "选择服务器",
                range(len(server_options)),
                format_func=lambda x: server_options[x]
            )
            
            selected_server = st.session_state.scan_results[selected_idx]
            
            # 操作按钮
            action_col1, action_col2, action_col3 = st.columns([1, 1, 2])
            with action_col1:
                if st.button("选择此服务器"):
                    st.session_state.selected_server = {
                        'host': selected_server['host'],
                        'port': selected_server['port']
                    }
                    st.rerun()
            
            with action_col2:
                if st.button("连接此服务器"):
                    with st.spinner("正在连接..."):
                        try:
                            if server_manager.connect(
                                selected_server['host'],
                                selected_server['port']
                            ):
                                st.success("连接成功！")
                                st.rerun()
                            else:
                                st.error("连接失败")
                        except Exception as e:
                            st.error(f"连接错误: {str(e)}")
        
        elif scan_button:  # 如果扫描后没有结果
            st.warning("未发现服务器")

def show_chat_page(client: OllamaClient, state_manager: StateManager):
    """显示聊天页面"""
    chat_ui = ChatInterface(client, state_manager)
    chat_ui.render()

def main():
    """主函数"""
    # 配置页面
    st.set_page_config(
        page_title="Ollama Chat",
        page_icon="💬",
        layout="wide"
    )
    
    # 添加代码高亮支持
    st.markdown("""
        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/styles/default.min.css">
        <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/highlight.min.js"></script>
        <script>hljs.highlightAll();</script>
        <style>
            pre { padding: 10px; border-radius: 4px; background-color: #f6f8fa; }
            code { font-family: 'Consolas', 'Monaco', monospace; }
        </style>
    """, unsafe_allow_html=True)
    
    # 配置日志
    logger = setup_logging()
    
    try:
        # 初始化组件
        if 'server_manager' not in st.session_state:
            st.session_state.server_manager = ServerManager()
            logger.info("初始化服务器管理器")
        
        # 检查服务器连接
        if not st.session_state.server_manager.get_current_server():
            show_connection_page(st.session_state.server_manager)
            return
        
        # 初始化客户端
        if 'client' not in st.session_state:
            server = st.session_state.server_manager.get_current_server()
            st.session_state.client = OllamaClient(
                host=server['host'],
                port=server['port']
            )
            logger.info(f"初始化API客户端: {server['host']}:{server['port']}")
        
        # 初始化状态管理器
        if 'state_manager' not in st.session_state:
            st.session_state.state_manager = StateManager()
            logger.info("初始化状态管理器")
        
        # 显示聊天界面
        show_chat_page(
            st.session_state.client,
            st.session_state.state_manager
        )
        
    except Exception as e:
        logger.error(f"应用运行错误: {str(e)}")
        st.error(f"发生错误: {str(e)}")
        ErrorHandler().handle_error(e)

if __name__ == "__main__":
    main() 